package learning.pcfg.inference;

import java.util.ArrayList;
import java.util.Enumeration;
import java.util.List;

import learning.pcfg.inference.Chart.ChartCell;
import learning.pcfg.inference.IndexedGrammar.Production;
import learning.pcfg.inference.IndexedGrammar.ProductionType;
import learning.pcfg.model.TerminalScore;

public class CYKParser {

	private IndexedGrammar grammar;
	private IParseScorer parseScorer;
	
	public CYKParser(IndexedGrammar grammar, IParseScorer parseScorer) {
		this.grammar = grammar;
		this.parseScorer = parseScorer;
	}
	
	public Parse parse(Region input) {
		
		// these can be set later to make inference more efficient
		int maxWordsInTerminalProduction = input.length();
		int maxWordsInInnerProduction = input.length();
		
		Chart chart = new Chart(input.length());
		
		Parse bestParse = null;
		double bestParseScore = 0;
		
		// terminal productions
		for (int numWords = 1; numWords <= maxWordsInTerminalProduction; numWords++) {
			Region subRegion = new Region(input.start, numWords);
			while (subRegion.end <= input.end)
			{
				ChartCell cc = chart.get(subRegion.start, subRegion.end);
				if (cc == null) cc = new ChartCell();
				
				applyTerminalProductions(subRegion, cc);
				
				if (cc.size() > 0) {
					chart.set(subRegion.start, subRegion.end, cc);
					
					// update best parse here, don't we only want to do that, if the region covers everything?
					Parse sp = cc.get(grammar.START_SYMBOL);
					if (sp != null) {
						double score = parseScorer.scoreParse(sp);
						if (bestParse == null || score > bestParseScore) {
							bestParse = sp;
							bestParseScore = score;
						}
					}
				}
				subRegion.start++;
				subRegion.end++;
			}
		}
		
		// inner productions
		for (int numWords = 1; numWords <= maxWordsInInnerProduction; numWords++) {
			Region subRegion = new Region(input.start, numWords);
			while (subRegion.end <= input.end)
			{
				ChartCell cc = chart.get(subRegion.start, subRegion.end);
				if (cc == null) cc = new ChartCell();
				
				applyBinaryProductions(subRegion, cc, chart);
				
				// apply unary productions, only if changes happened
				int prevSize = -1;
				while (prevSize < cc.size()) {
					prevSize = cc.size();
					applyUnaryProductions(subRegion, cc, chart);
				}
				
				if (cc.size() > 0) {
					chart.set(subRegion.start, subRegion.end, cc);
					
					// update best parse here, don't we only want to do that, if the region covers everything?
					Parse sp = cc.get(grammar.START_SYMBOL);
					if (sp != null) {
						double score = parseScorer.scoreParse(sp);
						if (bestParse == null || score > bestParseScore) {
							bestParse = sp;
							bestParseScore = score;
						}
					}
				}
				subRegion.start++;
				subRegion.end++;
			}
		}
		
		ChartCell cc = chart.get(input.start, input.end);
		if (cc != null)
			return cc.get(grammar.START_SYMBOL);
		else
			return null;
	}	
	
	private void applyTerminalProductions(Region region, ChartCell chartRegion) {
		List<TerminalScore> scoredTerminals = new ArrayList<TerminalScore>();
		parseScorer.scoreTerminals(region, scoredTerminals);
		for (TerminalScore ts : scoredTerminals) {
			Parse parse = new Parse(ts.production, null, region);
			parse.score = ts.score;
			chartRegion.set(ts.lhs, parse);
		}
	}
	
	private void applyUnaryProductions(Region region, ChartCell cc, Chart chart) {
		Enumeration<Production> e = grammar.productions(ProductionType.UNARY);
		while (e.hasMoreElements()) {
			Production prod = e.nextElement();

            Parse rhsParse = cc.get(prod.rhs1);
            if (rhsParse != null)
            {
                Parse lhsParse = new Parse(prod.id, new Parse[] { rhsParse }, region);
                lhsParse.score = parseScorer.scoreParse(lhsParse);

                Parse oldParse = cc.get(prod.lhs);

                if (oldParse == null || oldParse.score < lhsParse.score)
                {
                    cc.set(prod.lhs, lhsParse);
                }
            }
        }
	}
	
	private void applyBinaryProductions(Region region, ChartCell cc, Chart chart) {
        // enumerate all partitions of the input into two non-empty regions
        Region r1 = new Region(region.start, -1), r2 = new Region(-1, region.end);
        for (int breakPos = region.start + 1; breakPos < region.end; breakPos++)
        {
            r1.end = r2.start = breakPos;

            // lookup chart entries for the two regions
            ChartCell cr1 = chart.get(r1.start, r1.end);
            ChartCell cr2 = chart.get(r2.start, r2.end);

            // assuming that the two boxes in the chart are not empty
            if (cr1 != null && cr2 != null)
            {
            	Enumeration<Production> e = grammar.productions(ProductionType.BINARY);
            	while (e.hasMoreElements())
                {
            		Production prod = e.nextElement();

                    Parse rhs1parse = cr1.get(prod.rhs1);
                    Parse rhs2parse = cr2.get(prod.rhs2);

                    if (rhs1parse != null && rhs2parse != null)
                    {
                        Parse newParse = new Parse
                            (prod.id, new Parse[] { rhs1parse, rhs2parse }, region);
                        newParse.score = parseScorer.scoreParse(newParse);

                        Parse oldParse = cc.get(prod.lhs);

                        if (oldParse == null || oldParse.score < newParse.score)
                        {
                            cc.set(prod.lhs, newParse);
                        }
                    }
                }
            }
        }
	}
}
